"""
tellopy sample using keyboard and video player

Requires mplayer to record/save video.


Controls:
- tab to lift off
- WASD to move the drone
- space/shift to ascend/descent slowly
- Q/E to yaw slowly
- arrow keys to ascend, descend, or yaw quickly
- backspace to land, or P to palm-land
- enter to take a picture
- R to start recording video, R again to stop recording
  (video and photos will be saved to a timestamped file in ~/Pictures/)
- Z to toggle camera zoom state
  (zoomed-in widescreen or high FOV 4:3)
"""

import time
import pygame
import pygame.display
import pygame.key
import pygame.locals
import pygame.font
import os
import datetime
from subprocess import Popen, PIPE
from .MyTello import MyTello
import cv2
# from tellopy import logger

# log = tellopy.logger.Logger('TelloUI')




class TellKeyboard(object):
    """键盘控制无人机操作类 传入无人机对象"""
    def __init__(self,Tello:MyTello) -> None:
        self.speed = 20
        self.speed2 = 10
        self.drone = Tello
        self.controls = {
            'w': 'forward',
            's': 'backward',
            'a': 'left',
            'd': 'right',
            'space': 'up',
            'left shift': 'down',
            'right shift': 'down',
            'q': 'counter_clockwise',
            'e': 'clockwise',
            # arrow keys for fast turns and altitude adjustments
            'left': lambda drone, speed: drone.counter_clockwise(speed*2),
            'right': lambda drone, speed: drone.clockwise(speed*2),
            'up': lambda drone, speed: drone.up(speed*2),
            'down': lambda drone, speed: drone.down(speed*2),
            'tab': lambda drone, speed: drone.takeoff(),
            'backspace': lambda drone, speed: drone.land()
        }  
        
        pygame.init()
        pygame.display.init()
        self.screen=pygame.display.set_mode((640,480),0,32)

        self.prev_flight_data = None
        self.prev_info_data = None
        self.prev_battery_data = None
        self.drone.subscribe(self.drone.EVENT_FLIGHT_DATA,
                             self.flight_data_handler)
        self.drone.subscribe(self.drone.EVENT_BATTERY_DATA,
                             self.flight_battery_handler)
        self.drone.subscribe(self.drone.EVENT_INFO_DATA,
                             self.flight_info_handler)
        



    def keyController(self,frame):
        """通过键盘控制无人机,并显示无人机拍到视频图像
            'w'          : 'forward',
            's'          : 'backward',
            'a'          : 'left',
            'd'          : 'right',
            'space'      : 'up',
            'left shift' : 'down',
            'right shift': 'down',
            'q'          : 'counter_clockwise',
            'e'          : 'clockwise',
            
            用光标键控制飞行
            'left'       : counter_clockwise
            'right'      : clockwise
            'up'         : up
            'down'       : down
            'tab'        : takeoff
            'backspace'  : land        
        

        Args:
            frame (_mat_): 无人机视频图像

        Returns:
            _bool_: 当接收到ESC键则返回真（True），没接收到则返回假（False）
        """

        break_while=False
        try:
            frame=frame.transpose(1,0,2)
            frame=cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
            frame = pygame.surfarray.make_surface(frame)
            self.screen.blit(frame, (0,0))

            pygame.display.update()

            time.sleep(0.01)  # loop with pygame.event.get() is too mush tight w/o some sleep
            
            for e in pygame.event.get():
                # WASD for movement
                if e.type == pygame.locals.KEYDOWN:
                    print('+' + pygame.key.name(e.key))
                    keyname = pygame.key.name(e.key)
                    if keyname == 'escape':
                        # self.drone.close()
                        break_while=True
                        pygame.quit()
                        break
                    if keyname in self.controls:
                        key_handler = self.controls[keyname]
                        if type(key_handler) == str:
                            getattr(self.drone, key_handler)(self.speed)
                        else:
                            key_handler(self.drone, self.speed)

                elif e.type == pygame.locals.KEYUP:
                    print('-' + pygame.key.name(e.key))
                    keyname = pygame.key.name(e.key)
                    if keyname in self.controls:
                        key_handler = self.controls[keyname]
                        if type(key_handler) == str:
                            getattr(self.drone, key_handler)(0)
                        else:
                            key_handler(self.drone, 0)
            return break_while 
        except Exception as  e:
            print(str(e))
            pygame.quit()
            return True
        
    def write_hud(self, frame):
        """在图像上显示无人机信息"""
        if self.prev_flight_data == None:
            return frame
        stats = self.prev_flight_data.split('|')
        stats.extend( self.prev_info_data.split('|'))
        stats.append(self.prev_battery_data)
        # stats.append("Tracking:" + str(self.tracking))
        for idx, stat in enumerate(stats):
            text = stat.lstrip()
            cv2.putText(frame, text, (0, 30 + (idx * 30)),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.5, (0, 255, 0), lineType=30)
        return frame

    def flight_data_handler(self,  data):
        """Listener to flight data from the drone."""
        yaw, pitch, roll = data
        text = "Drone attitude:| yaw:{0},| pitch:{1},| roll:{2} ".format(
            yaw, pitch, roll)
        # print (text)
        if self.prev_flight_data != text:
            self.prev_flight_data = text
    def flight_info_handler(self,  data):
        """Listener to flight data from the drone."""
        alt, temp, motor_time = data
        text = "alt:{0},| temp:{1},| time:{2} ".format(
            alt, temp, motor_time)
        # print (text)
        if self.prev_info_data != text:
            self.prev_info_data = text
            
    def flight_battery_handler(self,  data):
        """Listener to flight data from the drone."""
        
        text = "battery:{0}".format(
            data)
        # print (text)
        if self.prev_battery_data != text:
            self.prev_battery_data = text            
        

            

if __name__ == '__main__':
    tello=MyTello('192.168.10.2')
    tellKeyboard=TellKeyboard(tello)
    tello.flight.motor_on()
    while True:
        frame=tello.get_video_frame()
        if tellKeyboard.keyController(frame):
            break
    tello.flight.motor_off()
    tello.quit()
    del tellKeyboard
    del tello
